﻿using Assets.Scripts.Bvh.Common;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using UnityEngine;

namespace Assets.Scripts.Bvh.Gpu
{
    class BvhGpuBuilder
    {
        public ComputeShader meshLoaderCs;
        public ComputeShader mortonCoderCs;
        public ComputeShader bitonicSortCs;
        public ComputeShader radixTreeBuilderCs;
        public ComputeShader aabbBuilderCs;

        private ComputeBuffer trianglesBuffer;
        private ComputeBuffer trianglesAABBBuffer;
        private ComputeBuffer leafNodesBuffer;
        private ComputeBuffer finalNodesBuffer;

        private RadixTreeNode[] LBvh;
        private List<Triangle> triangles;

        public BvhGpuBuilder()
        {
            meshLoaderCs = Resources.Load<ComputeShader>("ComputeShader/MeshLoaderCs");
            if (meshLoaderCs == null)
            {
                Debug.LogError("Failed to load ComputeShader MeshLoaderCs");
            }

            mortonCoderCs = Resources.Load<ComputeShader>("ComputeShader/MortonCoderCs");
            if (mortonCoderCs == null)
            {
                Debug.LogError("Failed to load ComputeShader MortonCoderCs");
            }

            bitonicSortCs = Resources.Load<ComputeShader>("ComputeShader/BitonicSortCs");
            if (bitonicSortCs == null)
            {
                Debug.LogError("Failed to load ComputeShader BitonicSortCs");
            }

            radixTreeBuilderCs = Resources.Load<ComputeShader>("ComputeShader/RadixTreeBuilderCs");
            if (radixTreeBuilderCs == null)
            {
                Debug.LogError("Failed to load ComputeShader RadixTreeBuilderCs");
            }

            aabbBuilderCs = Resources.Load<ComputeShader>("ComputeShader/AABBBuilderCs");
            if (aabbBuilderCs == null)
            {
                Debug.LogError("Failed to load ComputeShader AABBBuilderCs");
            }
        }

        public void Build(Mesh targetMesh)
        {
            var trianglesCount = LoadTriangles(targetMesh, out var meshMin, out var meshMax);
            PrepareData(trianglesCount, meshMin, meshMax);
            BitonicSort(trianglesCount);
            //DebugBuffer<RadixTreeNode>(leafNodesBuffer, trianglesCount);

            var interiorNodesBuffer = BuildRadixTree(trianglesCount);
            //DebugBuffer<RadixTreeNode>(interiorNodesBuffer, trianglesCount - 1);

            finalNodesBuffer = CollectNodes(trianglesCount, interiorNodesBuffer);
            //DebugBuffer<RadixTreeNode>(finalNodesBuffer, 2 * trianglesCount - 1);

            BuildAABB(trianglesCount, finalNodesBuffer);
            LBvh = FetchResult(2 * trianglesCount - 1, finalNodesBuffer);
        }

        public void Dispose()
        {
            trianglesAABBBuffer?.Dispose();
            leafNodesBuffer?.Dispose();
            finalNodesBuffer?.Dispose();
        }

        public void DrawGizmos()
        {
            if (LBvh == null) return;
            foreach (var node in LBvh)
            {
                var aabb = node.aabb;
                Gizmos.DrawWireCube(aabb.center, aabb.size);
            }
        }

        private int LoadTriangles(Mesh targetMesh, out Vector3 meshMinBound, out Vector3 meshMaxBound)
        {
            var bounds = targetMesh.bounds;
            meshMinBound = bounds.min;
            meshMaxBound = bounds.max;

            var meshLoaderKernel = meshLoaderCs.FindKernel("MeshLoader");

            var mesh = targetMesh;
            mesh.vertexBufferTarget |= GraphicsBuffer.Target.Raw;
            mesh.indexBufferTarget |= GraphicsBuffer.Target.Raw;

            var vertexBuffer = mesh.GetVertexBuffer(0);
            var indexBuffer = mesh.GetIndexBuffer();

            meshLoaderCs.SetBuffer(meshLoaderKernel, "_VertexBuffer", vertexBuffer);
            meshLoaderCs.SetBuffer(meshLoaderKernel, "_IndexBuffer", indexBuffer);

            var meshStride = mesh.GetVertexBufferStride(0);
            var indicesCount = (int)mesh.GetIndexCount(0);
            var verticeCount = mesh.vertexCount;

            meshLoaderCs.SetInt("_Stride", meshStride);
            meshLoaderCs.SetInt("_IndicesCount", indicesCount);
            meshLoaderCs.SetInt("_VerticesCount", verticeCount);

            var trianglesCount = indicesCount / 3;
            trianglesAABBBuffer = new ComputeBuffer(trianglesCount, Marshal.SizeOf<AABB>());
            meshLoaderCs.SetBuffer(meshLoaderKernel, "_TriangleAABBs", trianglesAABBBuffer);

            var dispatchX = trianglesCount / 1024 + 1;

            meshLoaderCs.Dispatch(meshLoaderKernel, dispatchX, 1, 1);

            vertexBuffer.Dispose();
            indexBuffer.Dispose();

            return trianglesCount;
        }

        private void PrepareData(int dataCount, Vector3 meshMin, Vector3 meshMax)
        {
            var mortonCoderkernel = mortonCoderCs.FindKernel("MortonCoder");
            mortonCoderCs.SetInt("_DataCount", dataCount);
            mortonCoderCs.SetVector("_MeshMinBound", meshMin);
            mortonCoderCs.SetVector("_MeshMaxBound", meshMax);

            leafNodesBuffer = new ComputeBuffer(dataCount, Marshal.SizeOf<RadixTreeNode>());

            mortonCoderCs.SetBuffer(mortonCoderkernel, "_TriangleAABBs", trianglesAABBBuffer);
            mortonCoderCs.SetBuffer(mortonCoderkernel, "_LeafNodes", leafNodesBuffer);

            var dispatchX = dataCount / 1024 + 1;
            mortonCoderCs.Dispatch(mortonCoderkernel, dispatchX, 1, 1);
        }

        private void BitonicSort(int rawCount)
        {
            var expandKernel = bitonicSortCs.FindKernel("ExpandData");
            var targetSortBuffer = leafNodesBuffer;
            var expanded = false;
            var dataCount = rawCount;

            //判断长度
            if (!Mathf.IsPowerOfTwo(dataCount))
            {
                var expandLength = Mathf.NextPowerOfTwo(dataCount);
                Debug.Log("Current data count : " + dataCount + "__expand length : " + expandLength);
                var expandBuffer = new ComputeBuffer(expandLength, Marshal.SizeOf<RadixTreeNode>());

                bitonicSortCs.SetInt("_DataLength", dataCount);
                bitonicSortCs.SetInt("_ExpandLength", expandLength);
                bitonicSortCs.SetBuffer(expandKernel, "_ExpandNodesBuffer", expandBuffer);
                bitonicSortCs.SetBuffer(expandKernel, "_LeafNodes", leafNodesBuffer);

                var dispatch = Mathf.CeilToInt(expandLength / 1024.0f);
                bitonicSortCs.Dispatch(expandKernel, dispatch, 1, 1);


                targetSortBuffer = expandBuffer;
                dataCount = expandLength;

                expanded = true;
            }

            var genSeqKernel = bitonicSortCs.FindKernel("GenerateSequence");
            var sortKernel = bitonicSortCs.FindKernel("BitonicSort");

            bitonicSortCs.SetBuffer(genSeqKernel, "_LeafNodes", targetSortBuffer);
            bitonicSortCs.SetInt("_DataLength", dataCount);
            var dispatchX = Mathf.CeilToInt(dataCount / 1024.0f);

            //双调序列生成
            {
                var signLength = 2;
                while (signLength < dataCount)
                {
                    var compareLength = signLength / 2;
                    bitonicSortCs.SetInt("_SignLength", signLength);

                    while (compareLength > 0)
                    {
                        bitonicSortCs.SetInt("_CompareLength", compareLength);
                        bitonicSortCs.Dispatch(genSeqKernel, dispatchX, 1, 1);
                        compareLength /= 2;
                    }
                    signLength *= 2;
                }
            }

            //双调排序
            {
                bitonicSortCs.SetBuffer(sortKernel, "_LeafNodes", targetSortBuffer);
                bitonicSortCs.SetInt("_DataLength", dataCount);

                var compareLength = dataCount / 2;
                while (compareLength != 0)
                {
                    bitonicSortCs.SetInt("_CompareLength", compareLength);
                    bitonicSortCs.Dispatch(sortKernel, dispatchX, 1, 1);

                    compareLength /= 2;
                }
            }

            //如果扩容了，再剔除一下。
            if (expanded)
            {
                var retrieveKernel = bitonicSortCs.FindKernel("RetrieveData");
                bitonicSortCs.SetInt("_DataLength", rawCount);
                bitonicSortCs.SetInt("_ExpandLength", dataCount);
                bitonicSortCs.SetBuffer(retrieveKernel, "_ExpandNodesBuffer", targetSortBuffer);
                bitonicSortCs.SetBuffer(retrieveKernel, "_LeafNodes", leafNodesBuffer);

                var dispatch = Mathf.CeilToInt(rawCount / 1024.0f);

                bitonicSortCs.Dispatch(retrieveKernel, dispatch, 1, 1);

                targetSortBuffer?.Dispose();
            }
        }

        private ComputeBuffer BuildRadixTree(int dataCount)
        {
            var interCount = dataCount - 1;
            var interiorNodesBuffer = new ComputeBuffer(interCount, Marshal.SizeOf<RadixTreeNode>());
            var tempInterior = new RadixTreeNode[interCount];
            for (int i = 0; i < tempInterior.Length; i++)
            {
                tempInterior[i] = new RadixTreeNode() { parentIndex = -1, nodeIndex = i, leftNodeIndex = -1, rightNodeIndex = -1 };
            }
            interiorNodesBuffer.SetData(tempInterior);

            var radixKernel = radixTreeBuilderCs.FindKernel("RadixTree");

            radixTreeBuilderCs.SetInt("_DataCount", dataCount);
            radixTreeBuilderCs.SetInt("_InteriorCount", interCount);
            radixTreeBuilderCs.SetBuffer(radixKernel, "_InteriorNodes", interiorNodesBuffer);
            radixTreeBuilderCs.SetBuffer(radixKernel, "_LeafNodes", leafNodesBuffer);

            var dispatchX = Mathf.CeilToInt(dataCount / 64.0f);
            radixTreeBuilderCs.Dispatch(radixKernel, dispatchX, 1, 1);

            return interiorNodesBuffer;
        }

        private ComputeBuffer CollectNodes(int dataCount, ComputeBuffer interiorNodesBuffer)
        {
            var collectKernel = radixTreeBuilderCs.FindKernel("CollectNodes");
            var allNodesCount = 2 * dataCount - 1;
            var allNodes = new ComputeBuffer(allNodesCount, Marshal.SizeOf<RadixTreeNode>());
            radixTreeBuilderCs.SetInt("_DataCount", dataCount);
            radixTreeBuilderCs.SetInt("_InteriorCount", dataCount - 1);

            radixTreeBuilderCs.SetBuffer(collectKernel, "_InteriorNodes", interiorNodesBuffer);
            radixTreeBuilderCs.SetBuffer(collectKernel, "_LeafNodes", leafNodesBuffer);
            radixTreeBuilderCs.SetBuffer(collectKernel, "_AllNodes", allNodes);

            var dispatchX = Mathf.CeilToInt(allNodesCount / 64.0f);
            radixTreeBuilderCs.Dispatch(collectKernel, dispatchX, 1, 1);

            interiorNodesBuffer.Dispose();

            return allNodes;
        }

        private void BuildAABB(int dataCount, ComputeBuffer allNodes)
        {
            var buildAABBKernel = aabbBuilderCs.FindKernel("BuildingAABB");
            aabbBuilderCs.SetInt("_DataCount", dataCount);
            aabbBuilderCs.SetInt("_InteriorCount", dataCount - 1);
            aabbBuilderCs.SetBuffer(buildAABBKernel, "_AllNodes", allNodes);

            var dispatchX = Mathf.CeilToInt(dataCount / 64.0f);
            aabbBuilderCs.Dispatch(buildAABBKernel, dispatchX, 1, 1);
        }

        private void DebugBuffer<T>(ComputeBuffer buffer, int count) where T : struct
        {
            var temp = new T[count];
            buffer.GetData(temp);
            foreach (var data in temp)
            {
                Debug.Log(data.ToString());
            }
        }

        private RadixTreeNode[] FetchResult(int dataCount, ComputeBuffer buffer)
        {
            var temp = new RadixTreeNode[dataCount];
            buffer.GetData(temp);
            return temp;
        }
    }
}
